from sqlalchemy import insert, select

from owl_admin.ext import db
from owl_ai.domain.rag_entity import KnowledgeBase
from owl_ai.domain.rag_po import KnowledgeBasePO
from owl_common.sqlalchemy.model import ColumnEntityList


class KnowledgeBaseMapper:
    default_fields = ['id', 'name', 'create_time', 'update_time', 'create_by', 'update_by']
    default_columns = ColumnEntityList(KnowledgeBasePO, default_fields, False)

    @classmethod
    def insert_one(cls, knowledge_base: KnowledgeBase):
        fields = ['id', 'name', 'create_time', 'update_time', 'create_by', 'update_by']
        data = knowledge_base.model_dump(include=fields, exclude_none=True, exclude_unset=True)
        stmt = insert(KnowledgeBasePO).values(data)
        out = db.session.execute(stmt).inserted_primary_key
        return out[0] if out else 0

    @classmethod
    def find_by_id(cls, knowledge_id: int):
        stmt = select(*cls.default_columns).where(KnowledgeBasePO.id == knowledge_id)
        row = db.session.execute(stmt).first()
        return cls.default_columns.cast(row, KnowledgeBase) if row else None
